#include "StdAfx.h"
#include "Picture.h"
#include "Utilities.h"
#include "Calculator.h"
#include "VectorSketch.h"
#include "Resource.h"

#include <fstream>
#include <sstream>
#include <iomanip>
#include <valarray>

#include <wincodecsdk.h>
#include <vector>
#include <algorithm>

#ifdef max
#undef max
#endif

// should be last include
#define STRSAFE_LIB
#include <strsafe.h>

const BITMAPINFOHEADER kGrayInfoTemplate = {
	sizeof(BITMAPINFOHEADER), 0, 0, 1, 8, BI_RGB,
	0, 3000, 3000, 256, 0};

double kNoiseLevel = 50;
double kNoiseLevelSquared = kNoiseLevel * kNoiseLevel;

const double kCriterion = 0.2;

const DWORD kUpdateIntervalInMilliseconds = 500;

#ifdef max
#undef max
#endif

Picture::Picture(Logger& logger)
	: mHandler(NULL),
	  mLogger(logger)
{
	mScale = 1;
	mSelectedViewID = kViewOriginal;
	mHighlightedStrokeIndex = kInvalidIndex;
	mBlur = 2;
	mThreshold = 0.1f;
	mNoise = 0.1f;
	mClutterSuppression = 10.0f;
	mIsEnabledSavePixels = false;
	mCurvatureCalculationToken = 0;

	// init grayinfo header
	mGrayInfo.i.bmiHeader = kGrayInfoTemplate;
	RGBQUAD* q = mGrayInfo.i.bmiColors;
	for (int i = 256; i--; ++q)
	{
		q->rgbRed = i;
		q->rgbGreen = i;
		q->rgbBlue = i;
	}
}

Picture::~Picture(void)
{
	Cleanup();
}

void Picture::SetEventHandler(EventHandler* handler)
{
	mHandler = handler;
}

ViewID Picture::GetView() const
{
	return mSelectedViewID;
}

BOOL Picture::Load(const TCHAR* path)
{
	BOOL result = false;
	
	Cleanup();

	Log(_T("LoadImage()"));
	HANDLE bmp = ::LoadImage(NULL, path, IMAGE_BITMAP, 0, 0, 
		LR_LOADFROMFILE | LR_CREATEDIBSECTION);
	if (NULL == bmp)
	{
		bmp = LoadUsingDecoders(path);
	}

	mBitmaps[kViewOriginal] = (HBITMAP)bmp;
	if (NULL == bmp) 
	{
		return false;
	}

	Log(_T("GetObject()"));
	BITMAP bitmap;
	if (!::GetObject(bmp, sizeof(bitmap), &bitmap))
	{
		return false;
	}
	
	int width = bitmap.bmWidth;
	int height = bitmap.bmHeight;
	mBitmap.SetWidthHeight(width, height);
	if (0 >= width || 0 >= height)
		return false;


	Log(_T("GetDIBits()"));
	int colorComponentsInRow = 3 * width;
	int bytesInRow = (colorComponentsInRow + 3) & (~3);

	BITMAPINFO bih = {{sizeof(BITMAPINFO), width, height,
		1, 24, BI_RGB, 0, 0, 0, 0, 0}, 0};

	int nL = ::GetDIBits(::GetDC(NULL), mBitmaps[kViewOriginal],
			0, height, mBitmap.GetData(), &bih, DIB_RGB_COLORS);


	mGrayInfo.i.bmiHeader.biWidth = width;
	mGrayInfo.i.bmiHeader.biHeight = height;

	Log(_T("UpdateDisplayLazily"));
	UpdateDisplayBitmapsLazily();

	return true;
}

BOOL Picture::SaveCurrentViewTo(const TCHAR* inPath) const
{
	HBITMAP hbmp = mBitmaps[mSelectedViewID];
	if (0 == hbmp) return false;

	DIBSECTION ds;
	if (!::GetObject(hbmp, sizeof(ds), &ds))
		return false;


	BITMAPFILEHEADER bfh;
	bfh.bfType = 'MB';
	bfh.bfReserved1 = bfh.bfReserved2 = 0;
	bfh.bfOffBits = sizeof(bfh) + ds.dsBmih.biSize + 
		ds.dsBmih.biClrUsed * sizeof(RGBQUAD);
	bfh.bfSize = bfh.bfOffBits + ds.dsBmih.biSizeImage;

	HANDLE hFile = ::CreateFile(inPath, GENERIC_WRITE, 0,
		NULL, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);
	if (INVALID_HANDLE_VALUE == hFile)
	{
		return false;
	}
	DWORD bytesWritten;
	::WriteFile(hFile, &bfh, sizeof(bfh), &bytesWritten, NULL);

	::WriteFile(hFile, &ds.dsBmih, sizeof(ds.dsBmih), &bytesWritten, NULL);
	// write palette
	if (int nPalette = ds.dsBmih.biClrUsed)
	{
		std::vector<RGBQUAD> palette(nPalette);

		HDC memDC = ::CreateCompatibleDC(0);
		::SelectObject(memDC, hbmp);
		::GetDIBColorTable(memDC, 0, nPalette, &palette[0]);
		::DeleteDC(memDC);

		::WriteFile(hFile, &palette[0], nPalette * sizeof(RGBQUAD), &bytesWritten, NULL);
	}
	::WriteFile(hFile, ds.dsBm.bmBits, ds.dsBmih.biSizeImage, &bytesWritten, NULL);
	::CloseHandle(hFile);
	return true;
}

void Picture::Draw(HDC ioHdc,const RECT& inRect) const
{
	if (0 == GetRawWidth() || 0 == GetRawHeight())
	{
		return;
	}

	if (mScale < 1 ||
		(mScale < 2 && mSelectedViewID == kViewC2Gradient))
	{
		::SetStretchBltMode(ioHdc, HALFTONE);
	}

	if (kViewBlank == mSelectedViewID)
	{
		RECT pictureRect;
		pictureRect.left = pictureRect.top = 0;
		pictureRect.right = GetWidth();
		pictureRect.bottom = GetHeight();
		
		RECT r;
		::IntersectRect(&r, &pictureRect, &inRect);

		HBRUSH whiteBrush = ::CreateSolidBrush(RGB(255, 255, 255));
		::FillRect(ioHdc, &r, whiteBrush);
		::DeleteObject(whiteBrush);
	}
	else
	{
		DrawBitmap(ioHdc, mBitmaps[mSelectedViewID], inRect);
	}
	DrawMargins(ioHdc, inRect);
	mSketch.DrawSketch(*this, ioHdc, inRect);
}

void Picture::GetDotShift(float& outShiftX, float& outShiftY) const
{
	switch (mSelectedViewID)
	{
		case kViewGradient:
		case kViewGradientNoBrightSpots:
		case kViewC2Gradient00:
			outShiftX = outShiftY = 0; break;
		case kViewC2Gradient: outShiftX = outShiftY = -0.25f; break;
		case kViewGradientSimple: outShiftX = outShiftY = 0.25f; break;
		case kViewC2Gradient01: outShiftX = -0.5f; outShiftY = 0; break;
		case kViewC2Gradient10: outShiftX = 0; outShiftY = -0.5f; break;
		case kViewC2Gradient11: outShiftX = outShiftY = -0.5; break;
		default: outShiftX = outShiftY = 0.5f; break;
	}
}

inline LONG downUpdate(LONG& ioDest, double inScale)
{
	double result = ::floor(ioDest / inScale);
	ioDest = LONG(result * inScale);
	return LONG(result);
}

inline LONG upUpdate(LONG& ioDest, double inScale)
{
	double result = ::ceil(ioDest / inScale);
	ioDest = LONG(result * inScale);
	return LONG(result);
}

void Picture::DrawBitmap(HDC ioHdc, HBITMAP inHBitmap, const RECT& inRect) const
{
	if (0 == inHBitmap) return;

	double scale = mScale;
	if (mSelectedViewID == kViewC2Gradient)
	{
		scale /= 2;
	}

	RECT sourceRect, destRect = inRect;
	sourceRect.left = downUpdate(destRect.left, scale);
	sourceRect.top = downUpdate(destRect.top, scale);
	sourceRect.right = upUpdate(destRect.right, scale);
	sourceRect.bottom = upUpdate(destRect.bottom, scale);

	HDC memDC = ::CreateCompatibleDC(ioHdc);
	HGDIOBJ oldObj = ::SelectObject(memDC, inHBitmap);

	::StretchBlt(ioHdc, destRect.left, destRect.top,
		destRect.right - destRect.left, destRect.bottom - destRect.top,
		memDC, sourceRect.left, sourceRect.top,
		sourceRect.right - sourceRect.left, sourceRect.bottom - sourceRect.top, SRCCOPY);

	::SelectObject(memDC, oldObj);
	::DeleteDC(memDC);
}

void Picture::DrawMargins(HDC ioHdc, const RECT& inUpdateRect) const
{
	LONG bitmapRight = GetWidth();
	LONG bitmapBottom = GetHeight();
	RECT rightMarginRect = inUpdateRect;
	RECT bottomMarginRect = inUpdateRect;
	bool shouldDrawRightMargin = false;
	bool shouldDrawBottomMargin = false;

	if (inUpdateRect.right > bitmapRight)
	{
		if (rightMarginRect.left < bitmapRight)
		{
			rightMarginRect.left = bitmapRight;
		}
		bottomMarginRect.right = bitmapRight;
		shouldDrawRightMargin = true;
	}
	if (inUpdateRect.bottom > bitmapBottom &&
		bottomMarginRect.left < bottomMarginRect.right)
	{
		if (bottomMarginRect.top < bitmapBottom)
		{
			bottomMarginRect.top = bitmapBottom;
		}
		shouldDrawBottomMargin = true;
	}
	if (!shouldDrawRightMargin && !shouldDrawBottomMargin)
	{
		return;
	}
	
	if (HBRUSH backgroundBrush = ::CreateSolidBrush(RGB(200, 200, 200)))
	{
		if (shouldDrawRightMargin)
		{
			::FillRect(ioHdc, &rightMarginRect, backgroundBrush);
		}
		if (shouldDrawBottomMargin)
		{
			::FillRect(ioHdc, &bottomMarginRect, backgroundBrush);
		}	
		::DeleteObject(backgroundBrush);
	}
}

void Picture::SetScale(float inScale)
{
	mScale = inScale;
}

void Picture::SelectViewAspect(ViewID inViewID)
{
	mSelectedViewID = inViewID;
	UpdateDisplayBitmapsLazily();
}

bool Picture::SetBlur(float inBlurRadius)
{
	if (mBlur == inBlurRadius)
	{
		return false;
	}
	mBlur = inBlurRadius;

	const int affectedViews[] = 
	{ kViewC2Gradient, kViewGradientSensitivity, kViewThreshold, 
		kViewC2Gradient00, kViewC2Gradient01, kViewC2Gradient10, kViewC2Gradient11, 0 };

	for (int i = 0; 0 != affectedViews[i]; ++i)
	{
		ClearBitmapAtIndex(affectedViews[i]);
	}
	UpdateDisplayBitmapsLazily();
	return true;
}

bool Picture::SetThreshold(float inThreshold)
{
	if (mThreshold != inThreshold)
	{
		mThreshold = inThreshold;
		ClearBitmapAtIndex(kViewThreshold);
		UpdateDisplayBitmapsLazily();
		return true;
	}
	
	return false;
}

bool Picture::SetNoise(float inNoise)
{
	if (mNoise != inNoise)
	{
		mNoise = inNoise;
		ClearBitmapAtIndex(kViewThreshold);
		ClearBitmapAtIndex(kViewGradientSensitivity);
		UpdateDisplayBitmapsLazily();
		return true;
	}
	
	return false;
}

bool Picture::SetClutterSuppression(float inClutterSuppression)
{
	if (mClutterSuppression != inClutterSuppression)
	{
		mClutterSuppression = inClutterSuppression;
		ClearBitmapAtIndex(kViewThreshold);
		UpdateDisplayBitmapsLazily();
		return true;
	}
	
	return false;
}

bool Picture::SetCurvatureCalculationToken(int inCurvatureCalculationToken)
{
	if (mCurvatureCalculationToken != inCurvatureCalculationToken)
	{
		mCurvatureCalculationToken = inCurvatureCalculationToken;
		return true;
	}
	else
		return false;
}

int Picture::GetWidth() const
{
	int width = mBitmap.GetRawWidth();	
	return int(width * mScale);
}

int Picture::GetHeight() const
{
	int height = mBitmap.GetRawHeight();
	return int(height * mScale);
}

bool Picture::CanSaveBitmap() const
{
	return  (NULL != mBitmaps[mSelectedViewID]);
}

/////////////////////////////////////////////////////////////////////////
// Private

void Picture::Run(ThreadRunner::Controller& controller)
{
	Log(_T("preparing iterator"));	
	if(controller.IsCancelled())
		return;
	ThresholdingIterator iterator(mBitmap, mBlur, mThreshold, mNoise, 
		mClutterSuppression);

	Calculator calculator(mBitmap, mBlur);
	calculator.SetMaximumDotSize(mBlur * 4);
	calculator.SetCurvatureCalculationToken(mCurvatureCalculationToken);
	iterator.SetSubpixelInfo(calculator.GetInfoBuffer());

	Log(_T("Iterate all"));
	mSketch.Set(GetRawWidth(), GetRawHeight());

	DWORD nextShowTime = ::GetTickCount() + kUpdateIntervalInMilliseconds;

	int x, y, n = 0, y0=0;
	VectorSketch& vectorSketch(mSketch.GetVectorSketch());
	while (iterator.Next(x, y))
	{
		if (y != y0)
		{
			y0 = y;
			if (nextShowTime < ::GetTickCount())
			{
				nextShowTime += kUpdateIntervalInMilliseconds;
				int percent = (y * 100) / GetRawHeight();
				std::wostringstream os;
				os << _T("Contour dots progress ");
				os << (TCHAR)((percent / 10) + '0');
				os << (TCHAR)((percent % 10) + '0');
				os << _T("%");
				mLogger.SetCurrentLogMessage(os.str());
				NotifyChange(WM_USER_UPDATESTATUSBAR);
			}
		}
		if (controller.IsCancelled())
			return;
		vectorSketch.TraceNewStroke(x, y, calculator);
	}
	vectorSketch.FinalizeStrokes(calculator);

	Log(_T("Placing the distinctive points"));
	vectorSketch.AimTheDistinctivePoints(5.0f * mBlur);

	vectorSketch.Finalize(calculator);

	Log(_T("Finished"));

	if (vectorSketch.IsDrawingAllStrokes())
	{
		NotifyChange(WM_USER_UPDATECLIENT);
	}
}

HBITMAP Picture::CreateBitmap()
{
	Array grayData;

	switch (mSelectedViewID)
	{
	case kViewGradient:  
		{
			::MakeGradient(mBitmap, grayData); 
			grayData = sqrt(grayData);
		}
		break;

	case kViewGradientSimple: 
		::MakeSimpleGradient(mBitmap, grayData); 
		grayData = sqrt(grayData);
		break;
	case kViewC2Gradient:  
		::MakeC2Gradient(mBitmap, mBlur, grayData);
		mGrayInfo.i.bmiHeader.biWidth *= 2;
		mGrayInfo.i.bmiHeader.biHeight *= 2;
		break;
	case kViewGradientNoBrightSpots: 
		{
			::MakeGradient(mBitmap, grayData); 
			float mg = grayData.max() / 4;
			for (int i = grayData.size(); i--;)
			{
				if (grayData[i] > mg)
				{
					grayData[i] = mg;
				}
			}
			grayData = std::sqrt(grayData);
		}
		break;
	case kViewGradientCross:
		::MakeGradientCross(mBitmap, grayData); 
		grayData = sqrt(grayData);
		break;
	case kViewGradientSensitivity:
		::MakeSensitivity(mBitmap, mBlur, mThreshold, mNoise, mClutterSuppression, grayData);
		break;
	case kViewThreshold:
		::MakeThreshold(mBitmap, mBlur, mThreshold, mNoise, mClutterSuppression, grayData);
		grayData = sqrt(grayData);
		break;
	//-----
	case kViewC2Gradient00:
		::MakeC2GradientX(mBitmap, mBlur, 0, grayData);
		break;
	case kViewC2Gradient01:
		::MakeC2GradientX(mBitmap, mBlur, 1, grayData);
		break;
	case kViewC2Gradient10:
		::MakeC2GradientX(mBitmap, mBlur, 2, grayData);
		break;
	case kViewC2Gradient11:
		::MakeC2GradientX(mBitmap, mBlur, 3, grayData);
		break;
	}

	float gMax = grayData.max();
	if (gMax == 0)
	{
		gMax = 1;
	}
	if (kViewThreshold == mSelectedViewID)
	{
		// the raw threshold data was compared with the square of gradient;
		// so we need to use the squre root of it for visualization
		grayData = (mThreshold * 255.0f / ::sqrt(gMax)) * sqrt(grayData);
	}
	else
	{
		// make the range of values 0..255 that is suitable for bitmap
		grayData *= 255.0f / gMax;
	}

	BYTE *q;
	HBITMAP hbmp = ::CreateDIBSection(0, &mGrayInfo.i, 
			DIB_RGB_COLORS, (void**)&q, NULL, 0);
	if (0 == hbmp) return false;

	int bytesInRow = (mGrayInfo.i.bmiHeader.biWidth + 3) & (~3);
	const float* p = &grayData[0];
	for (int y = 0; y < mGrayInfo.i.bmiHeader.biHeight; ++y)
	{
		for (int x = 0; x < mGrayInfo.i.bmiHeader.biWidth; ++x)
		{
			q[x] = int(p[x]);
		}
		q += bytesInRow;
		p += mGrayInfo.i.bmiHeader.biWidth;
	}
	return hbmp;
}

bool Picture::UpdateDisplayBitmapsLazily()
{
	bool result = false;

	int width = GetRawWidth();	
	int height = GetRawHeight();
	if (0 != mBitmaps[mSelectedViewID] || 0 >= width || 0 >= height) 
		return false;	
	mGrayInfo.i.bmiHeader.biWidth = width;
	mGrayInfo.i.bmiHeader.biHeight = height;

	if( (mSelectedViewID == kViewOriginal) ||
		(mSelectedViewID == kViewBlank) )
		return false;

	HBITMAP hbmp = CreateBitmap();
	mBitmaps[mSelectedViewID] = hbmp;
	result = (NULL != hbmp);		
	if (result)
	{
		NotifyChange(WM_USER_UPDATECLIENT);
	}

	return result;
}

void Picture::NotifyChange(int inNotificationID) const
{
	if(mHandler)
		mHandler->NotifyChange(inNotificationID);
}

void Picture::Cleanup()
{
	for (int i = kViewIDsCount; i--;)
	{
		ClearBitmapAtIndex(i);
	}

	mBitmap.CleanUp();
	mLogger.Clear();
	mSketch.ResetSketch();
}

void Picture::ClearBitmapAtIndex(int inIndex)
{
	if (HBITMAP& currentBitmapInArray = mBitmaps[inIndex])
	{
		::DeleteObject(currentBitmapInArray);
		currentBitmapInArray = 0;
	}
}

void Picture::Log(const TCHAR* inText)
{
	mLogger.Log(inText);
	NotifyChange(WM_USER_UPDATESTATUSBAR);
}

float Picture::GetScale() const
{
	return mScale;
}

HBITMAP Picture::LoadUsingDecoders(const TCHAR* inFilePath)
{
	HBITMAP hBitmap = NULL;
	::CoInitialize(NULL);

	IWICImagingFactory*  factory = NULL;
	HRESULT hr = ::CoCreateInstance(CLSID_WICImagingFactory,NULL,
		CLSCTX_INPROC_SERVER, IID_IWICImagingFactory, (LPVOID*)(&factory));
	if (S_OK != hr || NULL == factory)
	{
		return NULL;
	}

	IWICBitmapDecoder* decoder = NULL;
	hr = factory->CreateDecoderFromFilename(inFilePath, NULL, GENERIC_READ, 
	  WICDecodeMetadataCacheOnLoad, &decoder);
	if (S_OK == hr)
	{
		UINT frameCount = 0;
		if (S_OK == decoder->GetFrameCount(&frameCount))
		{
			// get the first available frame
			for (UINT i = 0; i < frameCount && NULL == hBitmap; ++i)
			{
				IWICBitmapFrameDecode *frameDecoder = NULL;
				if (S_OK == decoder->GetFrame(i, &frameDecoder))
				{
					hBitmap = CreateSRGBBitmapHandle(frameDecoder, factory);
					frameDecoder->Release();
				}
			}
		}
		decoder->Release();
	}	
	factory->Release();
	return hBitmap;
}

HBITMAP Picture::CreateSRGBBitmapHandle(IWICBitmapFrameDecode* inDecoder,
	IWICImagingFactory* inFactory)
{
	HBITMAP hBitmap = NULL;

	IWICFormatConverter *formatConverter = NULL;
	HRESULT hr = inFactory->CreateFormatConverter(&formatConverter);
	if (S_OK == hr)
	{
		hr = formatConverter->Initialize(inDecoder,
			GUID_WICPixelFormat24bppBGR, WICBitmapDitherTypeNone,
			NULL, 0.0f, WICBitmapPaletteTypeCustom);
		if (S_OK == hr)
		{
			hBitmap = CreateHBitmap(formatConverter);
		}
		formatConverter->Release();
	}

	return hBitmap;
}

HBITMAP Picture::CreateHBitmap(IWICBitmapSource* inSource)
{
	UINT width, height;
	if (S_OK != inSource->GetSize(&width, &height) || 
		0 == width || 0 == height)
	{
		return NULL;
	}

	BITMAPINFOHEADER bh;
	bh.biSize = sizeof(bh);
	bh.biWidth = width;
	bh.biHeight =  - int(height);
	bh.biPlanes = 1;
	bh.biBitCount = 24;
	bh.biCompression = BI_RGB;
	bh.biSizeImage = 0;
	bh.biXPelsPerMeter = 4000;
	bh.biYPelsPerMeter = 4000;
	bh.biClrUsed = 0;
	bh.biClrImportant = 0;
	
	void *bitmapBits = NULL;
	HBITMAP hBitmap = ::CreateDIBSection(GetDC(0), (BITMAPINFO*) &bh, 
		DIB_RGB_COLORS, &bitmapBits, NULL, 0);

	if (NULL != hBitmap && NULL != bitmapBits)
	{
		int bytesPerRow = (3 * width + 3) & (~3);
		HRESULT hr = inSource->CopyPixels(NULL, bytesPerRow, 
			bytesPerRow * height, (BYTE*)bitmapBits);
		if (S_OK != hr)
		{
			::DeleteObject(hBitmap);
			hBitmap = NULL;
		}
		hr = hr;
	}

	return hBitmap;
}

int Picture::GetRawWidth() const
{
	return mBitmap.GetRawWidth();
}

int Picture::GetRawHeight() const
{
	return mBitmap.GetRawHeight();
}

Sketch& Picture::GetSketch()
{
	return mSketch;
}

const Sketch& Picture::GetSketch() const
{
	return mSketch;
}

